#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-03-26

import socket
import time
import threading
from spirems.msg_helper import (encode_msg, decode_msg, get_all_msg_types, def_msg, check_msg,
                                index_msg_header, decode_msg_header, print_table)
import sys
import os
import argparse
import json
from queue import Queue
from spirems.subscriber import Subscriber
from spirems.publisher import Publisher
from spirems.client import Client
from spirems import sms2cvimg, sms2pcl, cvimg2sms
import cv2
import numpy as np


def read_pcd(pcd_path):
    header = {}
    with open(pcd_path, 'rb') as f:
        for line in f:
            line = line.decode('utf-8').strip()
            if line.startswith("DATA"):
                header["DATA"] = line.split()[1]
                break
            if line:
                key, val = line.split(' ', 1)
                header[key] = val

        width = int(header["WIDTH"])
        height = int(header["HEIGHT"])
        fields = header["FIELDS"].split()
        size = list(map(int, header["SIZE"].split()))
        count = list(map(int, header["COUNT"].split()))
        dtype = header["TYPE"].split()
        point_size = sum(s * c for s, c in zip(size, count))
        num_points = width * height

        data = f.read()
        if header["DATA"] == "binary":
            dt = np.dtype([(f, np.float32) for f in fields])
            arr = np.frombuffer(data, dtype=dt, count=num_points)
            return arr.view(np.float32).reshape(num_points, -1)
        else:
            raise NotImplementedError("only support binary")


def save_pcd_binary(filename, points, fields=['x', 'y', 'z']):
    """
    将NumPy数组保存为二进制格式的PCD文件
    :param filename: 文件名（如：'cloud.pcd'）
    :param points: NumPy数组，形状为(N, 3)，数据类型建议为float32
    """
    # print(fields)
    # print(points.shape)
    header = [
        "# .PCD v0.7 - Point Cloud Data file format",
        "VERSION 0.7",
        f"FIELDS {' '.join(fields)}",
        f"SIZE {' '.join('4' for f in fields)}",
        f"TYPE {' '.join('F' for f in fields)}",
        f"COUNT {' '.join('1' for f in fields)}",
        f"WIDTH {len(points)}",
        "HEIGHT 1",
        "VIEWPOINT 0 0 0 1 0 0 0",
        f"POINTS {len(points)}",
        f"DATA binary\n"
    ]
    header = '\n'.join(header)
    # print(header)
    with open(filename, 'wb') as f:
        f.write(header.encode('ascii'))
        assert points.dtype == np.float32
        f.write(np.ascontiguousarray(points, dtype=np.float32).tobytes())
        # np.savetxt(f, points, fmt='%.6f %.6f %.6f')


class SMSBagRecorder(threading.Thread):
    def __init__(
        self,
        record_urls: list,
        record_types: list,
        save_path: str = '',
        ip: str = '127.0.0.1',
        port: int = 9094
    ):
        threading.Thread.__init__(self)
        self.record_urls = record_urls
        self.record_types = record_types
        self.save_path = save_path
        self.ip = ip
        self.port = port

        self.queue = Queue()
        self.readers = []
        print("Subscribe:")
        for url_, type_ in zip(record_urls, record_types):
            reader_ = Subscriber(
                url_, 'std_msgs::Null', lambda msg, url=url_: self.queue.put({'msg': msg, 'url': url, 'raw': sms2cvimg(msg) if msg['type'] == 'memory_msgs::RawImage' else sms2pcl(msg)}) if msg['type'] in ['memory_msgs::RawImage', 'memory_msgs::PointCloud'] else self.queue.put({'msg': msg, 'url': url}),
                ip=ip, port=port
            )
            print('  ' + url_)
            self.readers.append(reader_)

        self.is_running = True
        self.start()

    def release(self):
        self.is_running = False
        self.queue.put(None)
        for reader_ in self.readers:
            reader_.kill()

    def run(self):
        save_name = os.path.join(self.save_path, "smsbag_" + time.strftime("%Y-%m-%d_%H-%M-%S"))
        print('SmsBag: ' + save_name)
        sub_name = ''
        sub_dir = ''
        jsonl_fp = None
        while self.is_running:
            try:
                msg = self.queue.get(block=True)
                if msg is None:
                    break

                time_min = time.strftime("%Y-%m-%d_%H-%M")
                if sub_name != time_min:
                    print("Storage pressure:", self.queue.qsize())
                    if jsonl_fp is not None:
                        jsonl_fp.close()
                    sub_name = time_min
                    sub_dir = os.path.join(save_name, sub_name)
                    os.makedirs(sub_dir)
                    jsonl_fn = os.path.join(sub_dir, time_min + '.jsonl')
                    jsonl_fp = open(jsonl_fn, "w", encoding="utf-8")

                msg['msg']['url'] = msg['url']
                if msg['msg']['type'] == 'memory_msgs::RawImage':
                    img = msg['raw']
                    img_fn = os.path.join(sub_dir, msg['url'].replace('/', '-'))
                    # print(img_fn)
                    if not os.path.exists(img_fn):
                        os.makedirs(img_fn)
                    img_n = str(int(time.time() * 1000)) + '.jpg'
                    msg['msg']['img_n'] = img_n
                    cv2.imwrite(os.path.join(img_fn, img_n), img)
                elif msg['msg']['type'] == 'memory_msgs::PointCloud':
                    pcl = msg['raw']
                    pcl_fn = os.path.join(sub_dir, msg['url'].replace('/', '-'))
                    # print(pcl_fn)
                    if not os.path.exists(pcl_fn):
                        os.makedirs(pcl_fn)
                    pcd_n = str(int(time.time() * 1000)) + '.pcd'
                    msg['msg']['pcd_n'] = pcd_n
                    save_pcd_binary(os.path.join(pcl_fn, pcd_n), pcl, fields=msg['msg']['fields'])
                    # save_pcd_binary(os.path.join(pcl_fn, pcd_n), pcl[:, [1,2,3]], fields=['x', 'y', 'z'])

                jsonl_fp.write(json.dumps(msg['msg']) + '\n')
            except Exception as e:
                print(e)
                self.is_running = False
                if jsonl_fp is not None:
                    jsonl_fp.close()

        self.release()


def _record(
        topics, ip, port, 
        save_path=''
    ):
    client = Client(
        '/system/service',
        'std_msgs::String',
        'std_msgs::StringMultiArray',
        ip=ip,
        port=port,
        request_once=True
    )
    req = def_msg('std_msgs::String')
    req['data'] = 'topic list'
    results = client.request_once(req)
    if 'data' in results:
        if len(results['data']) > 1:
            available_topics = results['data'][1:]
            record_urls, record_types = [], []
            if topics is None:
                record_urls = [t[0] for t in available_topics]
                record_types = [t[1] for t in available_topics]
            else:
                for t1 in topics:
                    for t in available_topics:
                        if t1 == t[0]:
                            record_urls.append(t[0])
                            record_types.append(t[1])
            if len(record_urls) > 0:
                recorder = SMSBagRecorder(
                    record_urls, 
                    record_types, 
                    ip=ip, 
                    port=port, 
                    save_path=save_path
                )
                recorder.join()
            else:
                print("The topic you subscribed to does not exist.")
        else:
            print("There are no topics available to subscribe to.")


class SMSBagPlayer(threading.Thread):
    def __init__(
        self,
        smsbag_jsonls: list,
        ip: str = '127.0.0.1',
        port: int = 9094
    ):
        threading.Thread.__init__(self)
        self.smsbag_jsonls = smsbag_jsonls
        self.ip = ip
        self.port = port

        self._publishers = dict()
        self.queue = Queue(maxsize=1000)
        self.loading_thread = threading.Thread(target=self.loading)
        self.loading_thread.start()
        self.is_running = True
        self.start()

    def loading(self):
        remaining_data = []
        for i, smsbag_jsonl in enumerate(self.smsbag_jsonls):
            parsed_data = []
            with open(smsbag_jsonl, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        if data['type'].startswith('memory_msgs::'):
                            data['base_dir'] = str(os.path.dirname(smsbag_jsonl))
                        parsed_data.append(data)
                    except json.JSONDecodeError:
                        print("File: {}, Line: {}, Parsing failed.")

            parsed_data.extend(remaining_data)
            parsed_data.sort(key=lambda x: x["timestamp"])
            if i != len(self.smsbag_jsonls) - 1:
                rnum = int(len(parsed_data) / 2)
                remaining_data = parsed_data[-rnum:]
                parsed_data = parsed_data[:-rnum]
            else:
                remaining_data = []

            for data in parsed_data:
                dict_put = {'msg': data}
                data_exist = True
                if data['type'] == 'memory_msgs::RawImage':
                    img_fn = os.path.join(data['base_dir'], data['url'].replace('/', '-'), data['img_n'])
                    if os.path.exists(img_fn):
                        img_raw = cv2.imread(img_fn, cv2.IMREAD_UNCHANGED)
                        dict_put['raw'] = img_raw
                    else:
                        data_exist = False
                elif data['type'] == 'memory_msgs::PointCloud':
                    pcd_fn = os.path.join(data['base_dir'], data['url'].replace('/', '-'), data['pcd_n'])
                    if os.path.exists(pcd_fn):
                        points_np = read_pcd(pcd_fn)
                        dict_put['raw'] = points_np
                    else:
                        data_exist = False
                if data_exist:
                    self.queue.put(dict_put, block=True, timeout=None)

        while not self.queue.empty():
            time.sleep(1)
        self.release()

    def release(self):
        self.is_running = False
        self.queue.put(None)
        for url, pub in self._publishers.items():
            pub.kill()

    def run(self):
        timestamp_base = 0.0
        first_msg = True
        while self.is_running:
            msg = self.queue.get(block=True)
            if msg is None:
                break

            if first_msg:
                first_msg = False
                timestamp_base = msg['msg']['timestamp']
                timestamp_now = time.time()

            dt_now = time.time() - timestamp_now
            dt_base = msg['msg']['timestamp'] - timestamp_base
            while dt_now <= dt_base:
                time.sleep(0.001)
                dt_now = time.time() - timestamp_now
                dt_base = msg['msg']['timestamp'] - timestamp_base

            url = msg['msg']['url']
            if url not in self._publishers:
                self._publishers[url] = Publisher(url, msg['msg']['type'], self.ip, self.port)

            if msg['msg']['type'] == 'memory_msgs::RawImage':
                raw_msg = self._publishers[url].cvimg2sms_mem(msg['raw'], msg['msg']['frame_id'], msg['msg']['timestamp'])
                self._publishers[url].publish(raw_msg)
            elif msg['msg']['type'] == 'memory_msgs::PointCloud':
                raw_msg = self._publishers[url].pcl2sms_mem(msg['raw'], msg['msg']['fields'], msg['msg']['frame_id'], msg['msg']['timestamp'])
                self._publishers[url].publish(raw_msg)
            else:
                self._publishers[url].publish(msg['msg'])

        self.release()


def _play(
        smsbag_dir, ip, port
    ):
    if not os.path.isdir(smsbag_dir):
        sys.exit("The SmsBag path ({}) does not exist.".format(smsbag_dir))

    subdirs = sorted(
        [entry.name for entry in os.scandir(smsbag_dir) if entry.is_dir()],
        key=lambda x: x.lower()
    )
    smsbag_jsonls = []
    for subdir in subdirs:
        jsonl_path = os.path.join(smsbag_dir, subdir, subdir + '.jsonl')
        if os.path.exists(jsonl_path) and os.path.isfile(jsonl_path):
            smsbag_jsonls.append(jsonl_path)

    player = SMSBagPlayer(smsbag_jsonls, ip, port)
    player.join()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'cmd',
        nargs='+',
        help='Your Command (record, play).')
    parser.add_argument(
        '-s', '--save-path',
        type=str,
        default='',
        help='SmsBag saving path.')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP.')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port.')
    args = parser.parse_args()
    # print(args.ip)
    # print(args.port)
    # print(args.cmd)
    if len(args.save_path) > 0:
        if os.path.exists(args.save_path) and os.path.isdir(args.save_path):
            print("Save to {}.".format(args.save_path))
        else:
            sys.exit("The folder ({}) does not exist.".format(args.save_path))

    if args.cmd[0] in ['record', 'play']:
        if 'play' == args.cmd[0]:
            assert len(args.cmd) > 1, "Usage: smsbag play [smsbag_dir]"
            _play(args.cmd[1], args.ip, args.port)
        if 'record' == args.cmd[0]:
            topics = None
            if len(args.cmd) > 1:
                topics = args.cmd[1:]
            _record(topics, args.ip, args.port, save_path=args.save_path)
    else:
        print('[ERROR] Supported command, use record, play.')


if __name__ == '__main__':
    main()
